# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 23:12:06 2022

@author: 13503
"""

from torch.utils.data import Dataset
import torch


"""
@brief: This is a dataset class to load data from tensor 
"""
class TensorDataset(Dataset):
    def __init__(self, data:torch.tensor):
        super().__init__()
        self.data = data

    def __getitem__(self, idx:int):
        return idx, self.data[idx]

    def __len__(self):
        return self.data.shape[0]